import os
import time

from train import train
from train_tfrecord import train as train_record
from build_model_data_function import build_model_data_fn#,build_model_tfdataset
import json
import wandb
import numpy as np

# def tfrecord_run_and_return_test(config_path):
#     """Run the training and write submission .

#     Args:
#         group_name ([type]): [description]
#         train_name ([type]): [description]
#     """    


#     hyparam = json.load(open(config_path))


#     fold_path = hyparam["inputs_fold"]
#     output_path = hyparam["output_path"]
#     if not os.path.exists(output_path):
#         os.makedirs(output_path)

#     model_fn, data_fn, submission_fn = build_model_tfdataset(**hyparam)
#     wandb.init(project=hyparam["project_name"],group=hyparam["group_name"],name=hyparam["train_name"],config=hyparam,entity="trillionmonster")
    
#     model, val_probs, test_text, test_probs = train_record( model_fn, data_fn, **hyparam)

#     lines = submission_fn(test_text, test_probs)
    
#     open(os.path.join(output_path,"submission.txt"),mode="w").writelines(lines)

#     wandb.finish()

#     return test_probs

def run_and_return_test(config_path):
    """Run the training and write submission .

    Args:
        group_name ([type]): [description]
        train_name ([type]): [description]
    """    


    hyparam = json.load(open(config_path))


    fold_path = hyparam["inputs_fold"]
    output_path = hyparam["output_path"]
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    model_fn, data_fn, submission_fn = build_model_data_fn(**hyparam)
    wandb.init(project=hyparam["project_name"],group=hyparam["group_name"],name=hyparam["train_name"],config=hyparam,entity="trillionmonster")
    
    model, val_probs, test_text, test_probs = train( model_fn, data_fn, **hyparam)

    # np.save(os.path.join(output_path,"val_probs.npy"),val_probs)

    # np.save(os.path.join(output_path,"test_probs.npy"),test_probs)

    lines = submission_fn(test_text, test_probs)
    
    open(os.path.join(output_path,"submission.txt"),mode="w").writelines(lines)

    wandb.finish()

    return test_probs




def run(config_path):
    """Run the training and write submission .

    Args:
        group_name ([type]): [description]
        train_name ([type]): [description]
    """    


    hyparam = json.load(open(config_path))


    fold_path = hyparam["inputs_fold"]
    output_path = hyparam["output_path"]
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    model_fn, data_fn, submission_fn = build_model_data_fn(**hyparam)
    wandb.init(project=hyparam["project_name"],group=hyparam["group_name"],name=hyparam["train_name"],config=hyparam,entity="trillionmonster")
    
    model, val_probs, test_text, test_probs = train( model_fn, data_fn, **hyparam)

    np.save(os.path.join(output_path,"val_probs.npy"),val_probs)

    np.save(os.path.join(output_path,"test_probs.npy"),test_probs)

    lines = submission_fn(test_text, test_probs)
    
    open(os.path.join(output_path,"submission.txt"),mode="w").writelines(lines)

    wandb.finish()

    
    